import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import os
import time
from sklearn.metrics import mean_squared_error, r2_score
import joblib

# 导入模型模块
from random_forest_model import random_forest_pipeline
from xgboost_model import xgboost_pipeline
from lstm_model import lstm_pipeline

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'SimSun', 'Arial Unicode MS']  # 尝试多种中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

def run_model_comparison():
    """
    运行模型比较流程
    """
    print("开始模型比较流程...")
    
    # 存储结果
    results = {
        'model_name': [],
        'training_time': [],
        'mse': [],
        'rmse': [],
        'mae': [],
        'r2': []
    }
    
    # 1. 随机森林模型
    print("\n开始运行随机森林模型...")
    start_time = time.time()
    _, rf_metrics = random_forest_pipeline(tune_hyperparameters=False)
    rf_training_time = time.time() - start_time
    
    results['model_name'].append('随机森林')
    results['training_time'].append(rf_training_time)
    results['mse'].append(rf_metrics['mse'])
    results['rmse'].append(rf_metrics['rmse'])
    results['mae'].append(rf_metrics['mae'])
    results['r2'].append(rf_metrics['r2'])
    
    print(f"随机森林模型训练时间: {rf_training_time:.2f}秒")
    
    # 2. XGBoost模型
    print("\n开始运行XGBoost模型...")
    start_time = time.time()
    _, xgb_metrics = xgboost_pipeline(tune_hyperparameters=False)
    xgb_training_time = time.time() - start_time
    
    results['model_name'].append('XGBoost')
    results['training_time'].append(xgb_training_time)
    results['mse'].append(xgb_metrics['mse'])
    results['rmse'].append(xgb_metrics['rmse'])
    results['mae'].append(xgb_metrics['mae'])
    results['r2'].append(xgb_metrics['r2'])
    
    print(f"XGBoost模型训练时间: {xgb_training_time:.2f}秒")
    
    # 3. LSTM模型
    print("\n开始运行LSTM模型...")
    start_time = time.time()
    _, lstm_metrics = lstm_pipeline(epochs=100, batch_size=32, sequence_length=12)
    lstm_training_time = time.time() - start_time
    
    results['model_name'].append('LSTM')
    results['training_time'].append(lstm_training_time)
    results['mse'].append(lstm_metrics['mse'])
    results['rmse'].append(lstm_metrics['rmse'])
    results['mae'].append(lstm_metrics['mae'])
    results['r2'].append(lstm_metrics['r2'])
    
    print(f"LSTM模型训练时间: {lstm_training_time:.2f}秒")
    
    # 创建结果DataFrame
    results_df = pd.DataFrame(results)
    
    # 打印比较结果
    print("\n模型比较结果:")
    print(results_df.to_string(index=False))
    
    # 可视化比较结果
    visualize_model_comparison(results_df)
    
    # 保存比较结果
    save_comparison_results(results_df)
    
    print("\n模型比较流程完成!")
    return results_df

def visualize_model_comparison(results_df, output_dir='../static/model_comparison'):
    """
    可视化模型比较结果
    
    参数:
    results_df: 比较结果DataFrame
    output_dir: 输出目录
    """
    print("\n可视化模型比较结果...")
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 设置颜色
    colors = ['skyblue', 'lightgreen', 'salmon']
    
    # 1. R² 比较图
    plt.figure(figsize=(10, 6))
    bars = plt.bar(results_df['model_name'], results_df['r2'], color=colors)
    
    # 添加数值标签
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                 f'{height:.4f}', ha='center', va='bottom')
    
    plt.xlabel('模型')
    plt.ylabel('决定系数 (R²)')
    plt.title('各模型决定系数 (R²) 比较')
    plt.grid(True, alpha=0.3)
    plt.savefig(os.path.join(output_dir, '模型比较_R2.png'))
    plt.close()
    
    # 2. RMSE 比较图
    plt.figure(figsize=(10, 6))
    bars = plt.bar(results_df['model_name'], results_df['rmse'], color=colors)
    
    # 添加数值标签
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 0.005,
                 f'{height:.4f}', ha='center', va='bottom')
    
    plt.xlabel('模型')
    plt.ylabel('均方根误差 (RMSE)')
    plt.title('各模型均方根误差 (RMSE) 比较')
    plt.grid(True, alpha=0.3)
    plt.savefig(os.path.join(output_dir, '模型比较_RMSE.png'))
    plt.close()
    
    # 3. 训练时间比较图
    plt.figure(figsize=(10, 6))
    bars = plt.bar(results_df['model_name'], results_df['training_time'], color=colors)
    
    # 添加数值标签
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 0.5,
                 f'{height:.2f}s', ha='center', va='bottom')
    
    plt.xlabel('模型')
    plt.ylabel('训练时间 (秒)')
    plt.title('各模型训练时间比较')
    plt.grid(True, alpha=0.3)
    plt.savefig(os.path.join(output_dir, '模型比较_训练时间.png'))
    plt.close()
    
    # 4. 综合性能雷达图
    plt.figure(figsize=(12, 8))
    
    # 定义性能指标
    metrics = ['R²', 'RMSE', '训练时间']
    
    # 获取数据
    r2_values = results_df['r2'].tolist()
    rmse_values = results_df['rmse'].tolist()
    time_values = results_df['training_time'].tolist()
    
    # 标准化RMSE和训练时间 (值越小越好，所以取倒数并标准化)
    max_rmse = max(rmse_values)
    min_rmse = min(rmse_values)
    normalized_rmse = [1 - (val - min_rmse) / (max_rmse - min_rmse) if max_rmse != min_rmse else 0.5 for val in rmse_values]
    
    max_time = max(time_values)
    min_time = min(time_values)
    normalized_time = [1 - (val - min_time) / (max_time - min_time) if max_time != min_time else 0.5 for val in time_values]
    
    # 组合数据 - 决定系数无需标准化，已经在0-1之间
    model_data = {
        '随机森林': [r2_values[0], normalized_rmse[0], normalized_time[0]],
        'XGBoost': [r2_values[1], normalized_rmse[1], normalized_time[1]],
        'LSTM': [r2_values[2], normalized_rmse[2], normalized_time[2]]
    }
    
    # 计算角度
    angles = np.linspace(0, 2*np.pi, len(metrics), endpoint=False).tolist()
    angles += angles[:1]  # 闭合雷达图
    
    # 绘制雷达图
    ax = plt.subplot(111, polar=True)
    
    for model_name, values in model_data.items():
        values_closed = values + [values[0]]  # 闭合
        ax.plot(angles, values_closed, 'o-', linewidth=2, label=model_name)
        ax.fill(angles, values_closed, alpha=0.1)
    
    # 添加标签
    ax.set_thetagrids(np.array(angles[:-1]) * 180/np.pi, metrics)
    
    # 设置雷达图最大值和最小值
    ax.set_ylim(0, 1)
    
    # 添加图例
    plt.legend(loc='upper right', bbox_to_anchor=(0.1, 0.1))
    
    plt.title('各模型性能比较雷达图')
    plt.savefig(os.path.join(output_dir, '模型比较_雷达图.png'))
    plt.close()
    
    print(f"可视化结果已保存至: {output_dir}")

def save_comparison_results(results_df, output_dir='../static/model_comparison'):
    """
    保存比较结果
    
    参数:
    results_df: 比较结果DataFrame
    output_dir: 输出目录
    """
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 保存为CSV
    csv_path = os.path.join(output_dir, '模型比较结果.csv')
    results_df.to_csv(csv_path, index=False)
    
    # 保存为HTML
    html_path = os.path.join(output_dir, '模型比较结果.html')
    results_df.to_html(html_path, index=False)
    
    print(f"比较结果已保存至: {csv_path}")

if __name__ == "__main__":
    # 运行模型比较流程
    run_model_comparison()